package edu.zju.gis.dbfg.model.index;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.XZQUtil;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.util.HashMap;
import java.util.Map;

import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_CRS;
import static edu.zju.gis.dbfg.model.base.Constant.DEFAULT_PYRAMID_CONFIG;

public class FragmentIndex {
    private final static Logger logger = LoggerFactory.getLogger(FragmentIndex.class);

    public static void main(String[] args) throws Exception {
        //Setup Environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setMaster("local[4]");
        conf.setAppName("SHDI");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup Model Arguments
        logger.info("Setup arguments");
        FragmentIndexArgs mArgs = FragmentIndexArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(SHDIndex.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs()));
        CoordinateReferenceSystem targetCrs = DEFAULT_CRS();
        PyramidConfig targetPyramidConfig = DEFAULT_PYRAMID_CONFIG();

        MathTransform transform = null;
        if (!crs.equals(targetCrs)) {
            transform = CRS.findMathTransform(crs, targetCrs);
        }

        //Read data from datasource and map it to geoItem
        String input = mArgs.getInput();
        String type = mArgs.getType();
        String dlbm = mArgs.getDlbmIndex();
        String area = mArgs.getAreaIndex();
        String zldwdm = mArgs.getZldwdmIndex();
        String id = mArgs.getIdIndex();
        int level = XZQUtil.getXZQLevel(mArgs.getLevel());
        Map<String, Object> attrs = new HashMap<>();
        attrs.put("dlbm", dlbm);
        attrs.put("area", area);
        attrs.put("zldwdm", zldwdm);
        JavaRDD<GeoItem> gi = null;
        if (type.equals("ES")) {
            gi = DataReaderRDD.readToRDD2(type, jsc, input, attrs, null, null, transform);
        } else {
            gi = DataReaderRDD.readToRDD(type, jsc, input, id, attrs, transform);
        }

        Broadcast<Integer> levelbroad = jsc.broadcast(level);

        JavaPairRDD<String, Tuple2<Double, Integer>> zltbRDD = gi.mapToPair(new PairFunction<GeoItem, String, Tuple2<Double, Integer>>() {
            @Override
            public Tuple2<String, Tuple2<Double, Integer>> call(GeoItem input) throws Exception {
                String key = String.valueOf(input.getProp().get("zldwdm")).substring(0, levelbroad.getValue());
                Double darea = Double.valueOf(String.valueOf(input.getProp().get("area")));
                Tuple2<Double, Integer> value = new Tuple2<Double, Integer>(darea,1);
                return new Tuple2<String, Tuple2<Double, Integer>>(key, value);
            }
        }).reduceByKey(new Function2<Tuple2<Double, Integer>, Tuple2<Double, Integer>, Tuple2<Double, Integer>>() {
            @Override
            public Tuple2<Double, Integer> call(Tuple2<Double, Integer> input1, Tuple2<Double, Integer> input2) throws Exception {
                Double darea = input1._1+input2._1;
                Integer count = input1._2+input2._2;
                return new Tuple2<>(darea, count);
            }
        });

        Map<String, Tuple2<Double, Integer>> result = zltbRDD.collectAsMap();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        writer.write(String.format("%s,%s\n", zldwdm, "SHDI"));
        for (String item : result.keySet()) {
            Tuple2<Double, Integer> value = result.get(item);
            Double fravalue = value._2/value._1;
            String s = String.format("%s,%.9f", item, fravalue);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();
    }
}
